package org.jeecg.modules.agentsflex;

import com.agentsflex.core.llm.Llm;
import com.agentsflex.core.message.AiMessage;
import com.agentsflex.llm.qwen.QwenLlm;
import com.agentsflex.llm.qwen.QwenLlmConfig;
import com.alibaba.fastjson.JSONObject;
import io.swagger.annotations.Api;
import io.swagger.annotations.ApiOperation;
import lombok.extern.slf4j.Slf4j;
import org.jeecg.modules.agentsflex.utils.WeatherUtil;
import org.jeecg.modules.easyai.keywordCapture.service.IKeywordCaptureService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;

/**
 * <p>
 *ai问答
 * </p>
 *
 * @author 29shusheng
 * @date 2024-7-24 14:58:54
 */

@Api(tags="ai问答")
@Slf4j
@RestController
@RequestMapping( "/ai/agentsFlex")
public class AiAgentsFlex {

    @Autowired
    private IKeywordCaptureService keywordCaptureService;
    // 创建一个用来维护SSE连接的对象
    private SseEmitter sseEmitter;
    private static Map<String, SseEmitter> sseCache = new ConcurrentHashMap<>();


    @ApiOperation(value="ai回答", notes="ai回答")
    @GetMapping(value = "/chat")
    public void chat(@RequestParam(name="speak") String speak) throws Exception {
        QwenLlmConfig config = new QwenLlmConfig();
        // 申请地址 https://bailian.console.aliyun.com/?apiKey=1#/api-key
        config.setApiKey("sk-f07053***********");
        //可以选择不同的模型，默认是qwen-turbo
        config.setModel("qwen-turbo");

        Llm llm = new QwenLlm(config);
        //自定义问答实现
        if(speak.contains("天气")){
            /*FunctionPrompt prompt = new FunctionPrompt(speak, WeatherFunctions.class);
            FunctionMessageResponse response = llm.chat(prompt);
            Object result = response.getFunctionResult();*/
            JSONObject identification = keywordCaptureService.identification(speak);
            String areaName = identification.getString("time");
            String result = WeatherUtil.getWeatherInfo(String.join("", areaName));
            if (sseEmitter != null ) {
                try {
                    //message.getContent()实际的数据推送内容
                    sseEmitter.send(SseEmitter.event().name("message").data(result));
                } catch (IOException e) {
                    System.out.println("Error sending data to client: " + e.getMessage());
                    sseEmitter.completeWithError(e);
                }
            }
        }else {
            // 通用聊天
            //异步返回
            llm.chatStream(speak, (context, response) -> {
                AiMessage message = response.getMessage();
                if (sseEmitter != null ) {
                    try {
                        //message.getContent()实际的数据推送内容
                        sseEmitter.send(SseEmitter.event().name("message").data(message.getContent()));
                    } catch (IOException e) {
                        System.out.println("Error sending data to client: " + e.getMessage());
                        sseEmitter.completeWithError(e);
                    }
                }
            });
        }
    }


       @RequestMapping(value = "/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
       public SseEmitter stream() {
           // 设置超时时间，0表示不过期。默认30秒，超过时间未完成会抛出异常：AsyncRequestTimeoutException
           sseEmitter = new SseEmitter(60000L);
           String clientId = String.valueOf(UUID.randomUUID());
           sseCache.put(clientId, sseEmitter);
           //sseEmitter.complete();
           // 连接断开回调
           sseEmitter.onCompletion(() -> {
               sseCache.get(clientId).complete();
               sseCache.remove(clientId);
           });
           // 连接超时
           sseEmitter.onTimeout(()-> {
               sseCache.get(clientId).complete();
               sseCache.remove(clientId);
           });
           // 连接报错
           sseEmitter.onError((throwable) ->  {
               sseCache.get(clientId).complete();
               sseCache.remove(clientId);
           });
           return sseEmitter;
       }


}
